#include "numuCCMultiPiSelection.hxx"
#include "numuCCSelection.hxx"
#include "baseSelection.hxx"
#include "CutUtils.hxx"
#include "Parameters.hxx"
#include "SubDetId.hxx"

//*******************************************************************************
numuCCMultiPiSelection::numuCCMultiPiSelection(bool forceBreak): SelectionBase(forceBreak){
//*******************************************************************************

    //Define the steps.
    DefineSteps();
}

//*******************************************************************************
void numuCCMultiPiSelection::DefineSteps(){
//*******************************************************************************

    // Copy all steps from the numuCCSelection
    numuCCSelection numu;
    CopySteps(numu);

    //Additional actions for the multi-pi selection.
    AddStep(StepBase::kAction, "fill_summary", new FillSummaryAction_numuCCMultiPi());
    AddStep(StepBase::kAction, "find_pions", new FindPionsAction());
    AddStep(StepBase::kAction, "find_tpc1", new FindTPC1TracksAction());

    //Add a split to the trunk with 3 branches.
    AddSplit(3);

    //First branch is for CC-0pi
    AddStep(0, StepBase::kCut, "CC-0pi", new NoPionCut());

    //Second branch is for CC-1pi
    AddStep(1, StepBase::kCut, "CC-1pi", new OnePionCut());

    //Third branch is for CC-Other
    AddStep(2, StepBase::kCut, "CC-Other", new OthersCut());

    // Set the branch aliases to the three branches
    SetBranchAlias(0,"CC-0pi",0);
    SetBranchAlias(1,"CC-1pi",1);
    SetBranchAlias(2,"CC-Other",2);
}

//********************************************************************
bool numuCCMultiPiSelection::FillEventSummary(const AnaBoxB& box, AnaEventB& event, Int_t allCutsPassed[]){
//********************************************************************

    (void)box;
    //CC0pi
    if(allCutsPassed[0]){
        event.Summary->EventSample = nd280Samples::kFGD1NuMuCC0Pi;
    }  
    //CC1pi
    else if (allCutsPassed[1]){
        event.Summary->EventSample = nd280Samples::kFGD1NuMuCC1Pi;
    }
    //CCOther
    else if (allCutsPassed[2]){
        event.Summary->EventSample = nd280Samples::kFGD1NuMuCCOther;
    }

    // otherwise kUnassigned is used from the EventSummary constructor
    return (event.Summary->EventSample != nd280Samples::kUnassigned);
}


//*********************************************************************
bool FillSummaryAction_numuCCMultiPi::Apply(const AnaEventB& event, AnaBoxB& box) const{
//*********************************************************************

    if(!box.HMNtrack) return 1;

    event.Summary->LeptonCandidate[nd280Samples::kFGD1NuMuCC0Pi] = box.HMNtrack;
    event.Summary->LeptonCandidate[nd280Samples::kFGD1NuMuCC1Pi] = box.HMNtrack;
    event.Summary->LeptonCandidate[nd280Samples::kFGD1NuMuCCOther] = box.HMNtrack;

    for(int i = 0; i < 4; ++i){
        event.Summary->VertexPosition[nd280Samples::kFGD1NuMuCC0Pi][i] = box.HMNtrack->PositionStart[i];
        event.Summary->VertexPosition[nd280Samples::kFGD1NuMuCC1Pi][i] = box.HMNtrack->PositionStart[i];
        event.Summary->VertexPosition[nd280Samples::kFGD1NuMuCCOther][i] = box.HMNtrack->PositionStart[i];
    }
    if(box.HMNtrack->TrueTrack){
        event.Summary->TrueVertex[nd280Samples::kFGD1NuMuCC0Pi] = box.HMNtrack->TrueTrack->TrueVertex;
        event.Summary->TrueVertex[nd280Samples::kFGD1NuMuCC1Pi] = box.HMNtrack->TrueTrack->TrueVertex;
        event.Summary->TrueVertex[nd280Samples::kFGD1NuMuCCOther] = box.HMNtrack->TrueTrack->TrueVertex;
    }
    return 1;
}

//*********************************************************************
bool FindPionsAction::Apply(const AnaEventB& event, AnaBoxB& box) const{
//*********************************************************************

    SubDetId::SubDetEnum det = SubDetId::kFGD1;

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);

    if (useTPCPions) numuCCMultiPiUtils::FindTPCPions(event,box,det, useOldSecondaryPID); 
    if (useME) numuCCMultiPiUtils::FindMEPions(event,*ccmultipibox,det, prod5Cut);
    if (useFGDPions) numuCCMultiPiUtils::FindIsoFGDPions(event,*ccmultipibox,det);

    int nnegpions =  ccmultipibox->nNegativePionTPCtracks; 
    int npospions =  ccmultipibox->nPositivePionTPCtracks; 
    int isofgdpions = ccmultipibox->nIsoFGDPiontracks;  
    int michelelectrons = ccmultipibox->nMichelElectrons;  
    int npi0 = ccmultipibox->nPosPi0TPCtracks+ ccmultipibox->nElPi0TPCtracks;

    int pionFGD = michelelectrons;
    if (!michelelectrons && isofgdpions>0) pionFGD = 1;

    ccmultipibox->nPosPions = npospions+pionFGD;
    ccmultipibox->nOtherPions = nnegpions+npi0;

    return true;
}

//*********************************************************************
bool FindTPC1TracksAction::Apply(const AnaEventB&, AnaBoxB& box) const{
//*********************************************************************
    // Look for pions in positive tracks 

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);
    ccmultipibox->nTPC1tracks = 0;

    for(int i = 0; i < box.nPositiveTPCtracks; i++ ) {
        AnaTrackB *ptrack = box.PositiveTPCtracks[i];
        if( SubDetId::GetDetectorUsed(ptrack->Detector, SubDetId::kTPC1) ){
            ccmultipibox->TPC1tracks[ccmultipibox->nTPC1tracks] = ptrack; 
            ccmultipibox->nTPC1tracks++;
        }
    }

    // Look for pions in negative tracks 

    for(int i = 0; i < box.nNegativeTPCtracks; i++ ) {
        AnaTrackB *ptrack = box.NegativeTPCtracks[i];
        if( SubDetId::GetDetectorUsed(ptrack->Detector, SubDetId::kTPC1) ){
            ccmultipibox->TPC1tracks[ccmultipibox->nTPC1tracks] = ptrack; 
            ccmultipibox->nTPC1tracks++;
        }
    }
    return true;
}

//*********************************************************************
bool NoPionCut::Apply(const AnaEventB& event, AnaBoxB& box) const{
//*********************************************************************

    (void)event;

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);

    if( ccmultipibox->nPosPions + ccmultipibox->nOtherPions == 0 ) return true; 

    return false; 
}


//*********************************************************************
bool OnePionCut::Apply(const AnaEventB& event, AnaBoxB& box) const{
//*********************************************************************

    (void)event;

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);

    if( ccmultipibox->nOtherPions != 0 ) return false; 
    if( ccmultipibox->nPosPions == 1 ) return true; 

    return false; 
}


//*********************************************************************
bool OthersCut::Apply(const AnaEventB& event, AnaBoxB& box) const{
//*********************************************************************

    (void)event;

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);

    if( ccmultipibox->nOtherPions != 0 ) return true; 
    if( ccmultipibox->nPosPions > 1 ) return true; 

    return false; 
}

//*********************************************************************
void numuCCMultiPiUtils::FindTPCPions(const AnaEventB& event, AnaBoxB& box, SubDetId::SubDetEnum det, bool useOldSecondaryPID){
//*********************************************************************

    (void)event;

    AnaBoxCCMultiPi * ccmultipibox = static_cast<AnaBoxCCMultiPi*>(&box);

    ccmultipibox->nPositivePionTPCtracks = 0;
    ccmultipibox->nPosPi0TPCtracks = 0;

    // Look for pions in positive tracks 
    for(int i = 0; i < box.nPositiveTPCtracks; i++ ) {
        AnaTrackB *ptrack = box.PositiveTPCtracks[i];
        if (!cutUtils::FiducialCut(*ptrack,det)) continue;

        if(useOldSecondaryPID){
            if ( numuCCMultiPiUtils::TPCpionSelection(ptrack) ) {
                ccmultipibox->PositivePionTPCtracks[ccmultipibox->nPositivePionTPCtracks] = ptrack;
                ccmultipibox->nPositivePionTPCtracks++;
            }
            else if ( numuCCMultiPiUtils::TPCElPi0Selection(ptrack) ) {
                ccmultipibox->PosPi0TPCtracks[ccmultipibox->nPosPi0TPCtracks] = ptrack; 
                ccmultipibox->nPosPi0TPCtracks++;
            }
        }
        else {
            // For Positive tracks we distinguish pions, electrons and protons.
            double ElLklh = anaUtils::GetPIDLikelihood(*ptrack,1);  
            double ProtonLklh = anaUtils::GetPIDLikelihood(*ptrack,2); 
            double PionLklh = anaUtils::GetPIDLikelihood(*ptrack,3); 
            double norm = ElLklh+ProtonLklh+PionLklh;
            ProtonLklh /= norm; 
            ElLklh /= norm; 
            PionLklh /= norm; 

            if( ProtonLklh > ElLklh && ProtonLklh > PionLklh ) continue; // If the highest probability is a proton continue. 

            // Id associated to the largest of the two probabilities.
            if( PionLklh > ElLklh ){
                ccmultipibox->PositivePionTPCtracks[ccmultipibox->nPositivePionTPCtracks] = ptrack;
                ccmultipibox->nPositivePionTPCtracks++;
            }
            else {
                if( ptrack->Momentum > 900. ) continue; // This is mainly protons.
                ccmultipibox->PosPi0TPCtracks[ccmultipibox->nPosPi0TPCtracks] = ptrack; 
                ccmultipibox->nPosPi0TPCtracks++;
            }
        }
    }

    // Look for pions in negative tracks 
    ccmultipibox->nNegativePionTPCtracks = 0;
    ccmultipibox->nElPi0TPCtracks = 0;

    for(int i = 0; i <box.nNegativeTPCtracks; i++ ) {
        AnaTrackB *ptrack = box.NegativeTPCtracks[i];
        if (!cutUtils::FiducialCut(*ptrack,det)) continue;

        if( box.HMNtrack == ptrack ) continue; // Same as the leading track.

        if(useOldSecondaryPID) {
            if ( numuCCMultiPiUtils::TPCpionSelection(ptrack) ) {
                ccmultipibox->NegativePionTPCtracks[ccmultipibox->nNegativePionTPCtracks] = ptrack;
                ccmultipibox->nNegativePionTPCtracks++;
            } 
            else if (  numuCCMultiPiUtils::TPCElPi0Selection(ptrack) ) {
                ccmultipibox->ElPi0TPCtracks[ccmultipibox->nElPi0TPCtracks] = ptrack; 
                ccmultipibox->nElPi0TPCtracks++;
            }
        }
        else {
            // For Negative tracks we distinguish pions, electrons and protons.
            double ElLklh = anaUtils::GetPIDLikelihood(*ptrack,1);  
            double PionLklh = anaUtils::GetPIDLikelihood(*ptrack,3); 
            double norm = ElLklh+PionLklh;
            ElLklh /= norm; 
            PionLklh /= norm;

            if( PionLklh > 0.8 ){ // Id associated to the largest of the two probabilities.
                ccmultipibox->NegativePionTPCtracks[ccmultipibox->nNegativePionTPCtracks] = ptrack;
                ccmultipibox->nNegativePionTPCtracks++;
            }
            else{ 
                ccmultipibox->ElPi0TPCtracks[ccmultipibox->nElPi0TPCtracks] = ptrack; 
                ccmultipibox->nElPi0TPCtracks++;
            }
        }
    }
}


//*********************************************************************
void numuCCMultiPiUtils::FindIsoFGDPions(const AnaEventB& event, AnaBoxCCMultiPi& ccmultipibox, SubDetId::SubDetEnum det){
//*********************************************************************

    // Get all tracks with FGD and no TPC
    AnaTrackB* allFgdTracks[100];
    int nFgdTracks = anaUtils::GetAllTracksUsingFGDAndNoTPC(event, allFgdTracks);

    ccmultipibox.nIsoFGDPiontracks = 0;
    ccmultipibox.nIsoFGDElPi0tracks = 0;

    //loop over tracks
    for (int i=0;i<nFgdTracks;i++){
        AnaTrackB* track = allFgdTracks[i];
        // Apply the fiducial cut and PID
        if( numuCCMultiPiUtils::FGDpionSelection(track,det) ){ 
            ccmultipibox.IsoFGDPiontracks[ccmultipibox.nIsoFGDPiontracks] = track;
            ccmultipibox.nIsoFGDPiontracks++;
        }
        if( numuCCMultiPiUtils::FGDElPi0Selection(track,ccmultipibox,det) ){
            ccmultipibox.IsoFGDElPi0tracks[ccmultipibox.nIsoFGDElPi0tracks] = track;
            ccmultipibox.nIsoFGDElPi0tracks++;
        }
    }
}

//*********************************************************************
void numuCCMultiPiUtils::FindMEPions(const AnaEventB& event, AnaBoxCCMultiPi& ccmultipibox, SubDetId::SubDetEnum det, bool prod5Cut){
//*********************************************************************

    AnaFgdTimeBinB* arr[100];
    ccmultipibox.nMichelElectrons = anaUtils::GetFGDMichelElectrons(event,det,arr, prod5Cut);
    for(int i = 0; i < ccmultipibox.nMichelElectrons; ++i){
        ccmultipibox.MichelElectrons[i] = arr[i];
    }
}

//*********************************************************************
bool numuCCMultiPiUtils::TPCpionSelection(AnaTrackB *track){
//*********************************************************************
    if ( anaUtils::GetPIDLikelihood(*track,3) < 0.3  ) return false; 

    double cut1 = (anaUtils::GetPIDLikelihood(*track,0)+anaUtils::GetPIDLikelihood(*track,3))/(1.-anaUtils::GetPIDLikelihood(*track,2)); 

    if( track->Momentum < 500.  && cut1 < 0.8 ) return false;

    return true;

}
//*********************************************************************
bool numuCCMultiPiUtils::TPCElPi0Selection(AnaTrackB *track){
//*********************************************************************

    if( track->nTPCSegments == 0 ) return false; // There is no TPC segment

    bool seltrack = false; 

    if (!cutUtils::TrackQualityCut(*track)) return seltrack;

    if( track->Momentum < 50. ) return seltrack; 

    //  if( track->Charge > 0. && track->Momentum > 800.) return seltrack; 

    Float_t pulls[4];

    for(int i = 0; i < track->nTPCSegments ; i++ ){

        AnaTpcTrackB *tpcTrack = track->TPCSegments[i];

        if( !tpcTrack  ) continue; // Extra protection.

        // Pulls are: Muon, Electron, Proton, Pion
        anaUtils::ComputeTPCPull(*tpcTrack,pulls);

        if (TMath::Abs(pulls[0]) > 1.e+6 // muon pull
                || TMath::Abs(pulls[1]) > 1.e+6 // electron pull
                || TMath::Abs(pulls[2]) > 1.e+6 // proton pull
                || TMath::Abs(pulls[3]) > 1.e+6) continue; // pion pull

        if( pulls[1] < -2.0 || pulls[1] > 2.0  ) break;

        if( track->Charge > 0. &&  ( pulls[2] > -4.0 && pulls[2] < 8.0 )  ) break;

        seltrack = true; 

    }
    return seltrack;
}

//*********************************************************************
bool numuCCMultiPiUtils::FGDpionSelection(AnaTrackB *track, SubDetId::SubDetEnum det){
//*********************************************************************
    if( track->nFGDSegments != 1 ) return false; // There is no FGD segment or it has more than one FGD 

    AnaFgdTrackB *fgdTrack = track->FGDSegments[0];

    if( !fgdTrack  ) return false; // Extra protection. 

    if( TMath::Abs(fgdTrack->Pullp) > 1.e+6 || TMath::Abs(fgdTrack->Pullmu) > 1.e+6 || TMath::Abs(fgdTrack->Pullpi) > 1.e+6  ) return false; 

    // Check the pullnotset variable. 
    // This tells us if the pull was calculated by the algorithm. 
    // If it was not calculated this variable will be 1, so for calculated pulls fgdpulls!=1. 
    if(  fgdTrack->Pullno == 1 ) return false;

    // This tell us something about the geometry. 
    // If the algorithm says the particle started and stopped in 1st/ 2nd fgd this is equal 1 or 2. 
    //  If the particle stopped in 1st/2nd fgd, this is equal -1/-2. Other cases: 0.
    if( fgdTrack->Containment == 0 ) return false;
    if( det == SubDetId::kFGD1 && fgdTrack->Containment != 1 ) return false;
    if( det == SubDetId::kFGD2 && fgdTrack->Containment != 2 ) return false; 
    if( det == SubDetId::kFGD && fgdTrack->Containment != 1 && fgdTrack->Containment != 2) return false; 

    double cosFGDpion = fgdTrack->DirectionStart[2];

    if(cosFGDpion > -0.3 && cosFGDpion < 0.3) return false;/// only for systematics issues. Must be remove in the future!!

    if( fgdTrack->Pullpi < 2.5 && fgdTrack->Pullpi > -2.0 ) return true; 

    return false;


}

//*********************************************************************
bool numuCCMultiPiUtils::FGDElPi0Selection(AnaTrackB *track, AnaBoxCCMultiPi& ccmultipibox, SubDetId::SubDetEnum det){
//*********************************************************************

    if( track->nFGDSegments != 1 ) return false; // There is no FGD segment or it has more than one FGD 

    AnaFgdTrackB *fgdTrack = track->FGDSegments[0];

    if( !fgdTrack  ) return false; // Extra protection. 

    if( TMath::Abs(fgdTrack->Pullp) > 1.e+6 || TMath::Abs(fgdTrack->Pullmu) > 1.e+6 || TMath::Abs(fgdTrack->Pullpi) > 1.e+6  ) return false; 
    int michelelectrons = 0;

    // Check the pullnotset variable. 
    // This tells us if the pull was calculated by the algorithm. 
    // If it was not calculated this variable will be 1, so for calculated pulls fgdpulls!=1. 
    if(  fgdTrack->Pullno == 1 ) return false;

    // This tell us something about the geometry. 
    // If the algorithm says the particle started and stopped in 1st/ 2nd fgd this is equal 1 or 2. 
    //  If the particle stopped in 1st/2nd fgd, this is equal -1/-2. Other cases: 0.
    if( fgdTrack->Containment == 0 ) return false; 
    if( det == SubDetId::kFGD1 && fgdTrack->Containment != 1 ) return false; 
    if( det == SubDetId::kFGD2 && fgdTrack->Containment != 2 ) return false; 
    if( det == SubDetId::kFGD && fgdTrack->Containment != 1 && fgdTrack->Containment != 2) return false; 

    //michelelectrons = std_numuccbanff_actions::FindMEPions(det); 

    michelelectrons = ccmultipibox.nMichelElectrons;  

    // michelelectrons = std_actions::GetFGDMichelElectrons(det);

    if(michelelectrons > 0 && fgdTrack->Pullpi < -3.0 ) return true;

    if(michelelectrons == 0 && fgdTrack->Pullpi < -1.5 ) return true;


    return false;

}

